#!/usr/bin/env python3
"""
TMS MCP Server Demo
Demonstrates the full functionality of the TMS MCP Server
"""

import asyncio
import json
import subprocess
import sys
import time

class TMSMCPDemo:
    def __init__(self):
        self.process = None
        self.request_id = 0
    
    def get_next_id(self):
        self.request_id += 1
        return self.request_id
    
    async def start_server(self):
        """Start the MCP server"""
        print("🚀 Starting TMS MCP Server...")
        self.process = subprocess.Popen(
            ["./target/debug/tms-mcp", "--log-level", "info"],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            cwd="."
        )
        
        # Initialize the server
        init_request = {
            "jsonrpc": "2.0",
            "id": self.get_next_id(),
            "method": "initialize",
            "params": {
                "protocolVersion": "2025-06-18",
                "capabilities": {},
                "clientInfo": {
                    "name": "tms-demo-client",
                    "version": "1.0.0"
                }
            }
        }
        
        self.process.stdin.write(json.dumps(init_request) + "\n")
        self.process.stdin.flush()
        
        response_line = self.process.stdout.readline()
        response = json.loads(response_line.strip())
        
        if "result" not in response:
            raise Exception("Failed to initialize server")
        
        print("✅ Server initialized successfully")
        return True
    
    async def call_tool(self, tool_name, arguments=None):
        """Call a tool and return the response"""
        if arguments is None:
            arguments = {}
        
        request = {
            "jsonrpc": "2.0",
            "id": self.get_next_id(),
            "method": "tools/call",
            "params": {
                "name": tool_name,
                "arguments": arguments
            }
        }
        
        self.process.stdin.write(json.dumps(request) + "\n")
        self.process.stdin.flush()
        
        response_line = self.process.stdout.readline()
        response = json.loads(response_line.strip())
        
        return response
    
    async def demo_health_check(self):
        """Demo: Health check"""
        print("\n🏥 Health Check Demo")
        print("=" * 50)
        
        response = await self.call_tool("health_check")
        
        if "result" in response:
            content = response["result"]["content"][0]["text"]
            print(f"✅ Health Check Result: {content}")
        else:
            print(f"❌ Health check failed: {response.get('error', 'Unknown error')}")
    
    async def demo_system_stats(self):
        """Demo: System statistics"""
        print("\n📊 System Statistics Demo")
        print("=" * 50)
        
        response = await self.call_tool("get_system_stats")
        
        if "result" in response:
            content = response["result"]["content"][0]["text"]
            print(f"📈 System Stats: {content}")
        else:
            print(f"❌ System stats failed: {response.get('error', 'Unknown error')}")
    
    async def demo_list_vehicles(self):
        """Demo: List vehicles"""
        print("\n🚛 Vehicle Management Demo")
        print("=" * 50)
        
        print("📋 Listing existing vehicles...")
        response = await self.call_tool("list_vehicles", {
            "page": 1,
            "page_size": 5
        })
        
        if "result" in response:
            content = response["result"]["content"][0]["text"]
            print(f"🚗 Vehicles: {content}")
        else:
            print(f"❌ List vehicles failed: {response.get('error', 'Unknown error')}")
    
    async def demo_create_vehicle(self):
        """Demo: Create a new vehicle"""
        print("\n🆕 Creating a new vehicle...")
        
        # Generate a unique license plate
        timestamp = int(time.time())
        license_plate = f"DEMO-{timestamp % 10000}"
        
        response = await self.call_tool("create_vehicle", {
            "license_plate": license_plate,
            "model": "Mercedes Sprinter",
            "capacity": 3.5,
            "vehicle_type": "Van",
            "driver_id": "driver-001"
        })
        
        if "result" in response:
            content = response["result"]["content"][0]["text"]
            print(f"✅ Vehicle created: {content}")
            return license_plate
        else:
            print(f"❌ Create vehicle failed: {response.get('error', 'Unknown error')}")
            return None
    
    async def demo_get_vehicle(self, vehicle_id):
        """Demo: Get vehicle details"""
        print(f"\n🔍 Getting vehicle details for ID: {vehicle_id}")
        
        response = await self.call_tool("get_vehicle", {
            "id": vehicle_id
        })
        
        if "result" in response:
            content = response["result"]["content"][0]["text"]
            print(f"📄 Vehicle details: {content}")
        else:
            print(f"❌ Get vehicle failed: {response.get('error', 'Unknown error')}")
    
    async def demo_update_vehicle_status(self, vehicle_id):
        """Demo: Update vehicle status"""
        print(f"\n🔄 Updating vehicle status for ID: {vehicle_id}")
        
        response = await self.call_tool("update_vehicle_status", {
            "id": vehicle_id,
            "status": "In Transit",
            "notes": "Updated via MCP demo"
        })
        
        if "result" in response:
            content = response["result"]["content"][0]["text"]
            print(f"✅ Status updated: {content}")
        else:
            print(f"❌ Update status failed: {response.get('error', 'Unknown error')}")
    
    async def run_demo(self):
        """Run the complete demo"""
        try:
            await self.start_server()
            
            print("\n🎯 TMS MCP Server Feature Demo")
            print("=" * 60)
            
            # Demo 1: Health check
            await self.demo_health_check()
            
            # Demo 2: System statistics
            await self.demo_system_stats()
            
            # Demo 3: List vehicles
            await self.demo_list_vehicles()
            
            # Demo 4: Create a new vehicle
            new_vehicle_plate = await self.demo_create_vehicle()
            
            # Demo 5: List vehicles again to see the new one
            print("\n📋 Listing vehicles after creation...")
            await self.demo_list_vehicles()
            
            print("\n🎉 Demo completed successfully!")
            print("\n📝 Summary:")
            print("✅ MCP Server initialization")
            print("✅ Health check")
            print("✅ System statistics")
            print("✅ Vehicle listing")
            print("✅ Vehicle creation")
            print("✅ Real-time data integration with TMS backend")
            
        except Exception as e:
            print(f"❌ Demo failed: {e}")
            return False
        
        finally:
            if self.process:
                self.process.terminate()
                try:
                    self.process.wait(timeout=5)
                except subprocess.TimeoutExpired:
                    self.process.kill()
                    self.process.wait()
        
        return True

async def main():
    demo = TMSMCPDemo()
    success = await demo.run_demo()
    sys.exit(0 if success else 1)

if __name__ == "__main__":
    asyncio.run(main())